[None][feat] Update CuTeDSL MegaMoE kernels by Barry-Delaney · Pull Request #16190 · NVIDIA/TensorRT-LLM

Barry-Delaney · 2026-07-09T12:00:34Z

Dev Engineer Review

MegaMoE CuteDSL NVFP4 path: introduced lazy import + conditional custom-op registration for the MegaMoE kernel wrapper, plus major interface/codegen changes for tactic selection and token-aware execution.
- Replaced deterministic fallback constant with default_megamoe_tactic(num_tokens) and made candidate enumeration depend on num_tokens.
- Updated tactic representation to the new 8-tuple model; tightened validate_megamoe_tactic() arity/field constraints.
Runner/custom-op/kernel interface expansion:
- Added knobs: in_kernel_fc2_reduce, combine_format, tactic_autotune, and optional live num_tokens.
- Removed prior public knobs: token_back_by_dispatch, non_ubulk_fc2_store.
- Updated fake registration and call signatures to match the new public interface.
Kernel/scheduler correctness and determinism updates:
- Standardized token-validity fields in scheduler/work-tile flow (valid_tokens_in_cta_tile, and non-swap_ab cluster-tile + fc1_counter_index handling).
- CUDA-graph-safe workspace/fence lifecycle refactor (pointer-keyed fencing, weakref/ABA detection, guarded reuse).
- Added/updated quantization + combine plumbing (QuantImpl, combine_format propagation, centralized NVFP4/CombineFormat logic).
- Refactored token communication for quantized token-back using CombineFormat, TokenSrcMetadata, and GpuReleaseFlagBatchTracker; updated NVLink/barrier + scheduling behavior and tail reset delegation.
Additional MegaMoE CuteDSL infra changes:
- Removed import_topk_reduce export/function from mega_moe_nvfp4/__init__.py.
- Refactored epilogue routing (notably fc2 output routing via Fc2OutputRouter) and completion signaling via batch flag tracking.
- Added lazy constant resolution + reciprocal limit constants in megamoe_constants.py.
- Updated scheduler extension work-tile metadata field naming (valid_tokens_in_cta_tile).
Backend/integration:
- Added new MoE backend identifier MEGAMOE_CUTEDSL and wired it through model args/manifests/registries and benchmark harnesses.
- Enabled multi-rank adaptive max_tokens_per_rank bucketing.
- Tightened MPI/mpirun process-group bootstrapping invariants for EP/WORLD sizing.
- Implemented streamed-weight loading for NVFP4 MegaMoE CuteDSL; explicitly disabled hot weight reloading.
Tagged fix alignment:
- Commit TRTLLM-12950 enforces a minimum MegaMoE group hint (as called out in the provided commit summary).

Key checks before approval:

Verify all signature/API wiring is consistent end-to-end: runner ↔ Torch custom op (incl. fake registration) ↔ kernel, especially around combine_format, in_kernel_fc2_reduce, tactic_autotune, and num_tokens slicing.
Confirm tactic ordering/tuple layout is consistently used as the single source of truth across enumeration, validation, and autotune selection (including the default_megamoe_tactic(num_tokens) fallback path).
Sanity-check token-communication + token-back scheduling correctness for both non-tuning and autotune paths, including CUDA-graph reuse guards and workspace counter lifecycle.
For MEGAMOE_CUTEDSL, validate candidate filtering behavior and EP/WORLD invariants in distributed search paths.

CI status (from provided run notes):

Pipelines failed for commits 468c64a and ec63e67.
Commit a6aa781 failed in multiple runs (with and without fail-fast), then a later rerun succeeded.
Latest reported pipeline completed successfully.

QA Engineer Review

Test list changes

Modified (test-list / waive / QA-list related files):

tests/integration/test_lists/test-db/l0_b200.yml
- Added: unittest/_torch/modules/moe/test_megamoe_streaming_load.py
tests/integration/test_lists/qa/README.md
- Updated example pytest command to point at accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_auto_dtype
tests/integration/test_lists/test-db/README.md
- Updated specific accuracy/test cases listed in the README.
tests/integration/test_lists/waives.txt
- Added:
  - unittest/disaggregated/test_kv_transfer.py SKIP (nvbugs/6403793)
  - unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (nvbugs/6445316)
- Removed:
  - unittest/executor/test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release SKIP (nvbugs/6480574)
Also modified multiple other tests/integration/test_lists/test-db/l0_*.yml and README files in tests/integration/test_lists/test-db/ (mostly deletions/rewording per diff stats).

Verdict for test-list coverage: needs follow-up (test selection/list churn is broader than the new MegaMoE test entry, and additional coverage intent is unclear from provided context).

Test code changes (outside test lists)

New/modified pytest unit tests detected from diff (files under tests/unittest/**):

tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
- Added: test_initial_streaming_load_layer_atomic
- Added: test_partial_load_rejected_before_source_materialization
tests/unittest/_torch/modules/moe/test_moe_backend.py
- Added/updated:
  - test_enumerate_megamoe_candidate_tactics_curated_space
  - test_megamoe_cutedsl_tactic_autotune_defaults_off
  - test_megamoe_cutedsl_tuning_mode_forces_top_maxt_bucket
Other touched unit-test files with added tests (not indicated as added to tests/integration/test_lists/test-db/ in the provided diff):
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
  - test_msa_metadata_rejects_undersized_max_score_buffer
  - test_msa_proxy_max_score_view_is_contiguous_over_stable_store
  - test_msa_requires_block_size_128
  - test_resolver_selects_msa_backend_when_available
- tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py
  - test_aggregate_across_ranks_stays_within_available_memory
  - test_colocated_ranks_divide_node_memory_budget
  - test_memlock_limit_caps_quota
  - test_non_positive_result_falls_back_to_device_quota
  - test_result_is_always_positive
  - test_single_rank_uses_device_quota_when_memory_is_ample
  - test_unknown_limits_fall_back_to_device_quota
- tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py
  - test_gate_bias_params_map
  - test_kv_scale_is_not_expanded_when_kv_heads_cover_tp
  - test_load_weights_accepts_base_mapper_without_params_map
  - test_load_weights_exposes_weight_mapper
  - test_mapper_registration_and_mx_fallback
  - test_nvfp4_scale_behavior_is_preserved
  - test_quant_config_none_is_guarded
  - test_tp8_mxfp8_duplicates_kv_weight_and_scale_via_callbacks
- tests/unittest/_torch/models/test_minimax_m3.py
  - test_text_norm_weights_real_loader_smoke
- tests/unittest/_torch/sampler/test_torch_sampler.py
  - test_out_of_range_decay_params_rejected
  - test_reject_speculative_draft_tokens
  - test_runtime_update_parity
  - test_strategy_routing
- tests/unittest/_torch/speculative/test_eagle3.py
  - test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31
- tests/unittest/llmapi/apps/_test_openai_metrics.py
  - test_metrics_available_before_first_request
- tests/unittest/llmapi/apps/_test_openai_prometheus.py
  - test_kv_cache_metrics_available_before_first_request
- tests/unittest/tools/test_config_database_sync.py
  - test_profile_and_validated_commit_metadata
  - test_recipe_metadata_validation

Coverage / verdict:

Confirmed test-list CI coverage: test_megamoe_streaming_load.py is added to l0_b200 test list.
Other newly added/modified unit tests are not shown as added to tests/integration/test_lists/test-db/ or tests/integration/test_lists/qa/ in the provided diff.
Verdict: needs follow-up due to limited/unclear CI coverage for the broader set of unit-test changes, plus the earlier CI failure history requiring confirmation on the final successful rerun.

coderabbitai · 2026-07-22T05:43:22Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8a302a21-b63c-49ea-98a5-dc23360c25e6

📥 Commits

Reviewing files that changed from the base of the PR and between a6aa781 and f047416.

📒 Files selected for processing (2)

tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py
tests/unittest/_torch/modules/moe/test_moe_backend.py

🚧 Files skipped from review as they are similar to previous changes (2)

tests/unittest/_torch/modules/moe/test_moe_backend.py
tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py

Walkthrough

MegaMoE CuteDSL adds quantized combine formats, token-aware tactics, adaptive token buckets, CUDA-graph-safe workspace fencing, distributed token-back communication, device-side top-k reduction, and streamed NVFP4 weight loading. Backend registration, scheduler integration, model configuration, and tests are updated.

Changes

MegaMoE CuteDSL runtime

Layer / File(s)	Summary
Runtime contracts and tactic selection `tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py`, `tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/{contract.py,megamoe_constants.py}`	Custom-op arguments, eight-field tactics, token-aware enumeration, workspace caching, profiling scratch, fencing, lazy imports, and runtime contract evaluation are updated.
Kernel workspace and token communication `tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/{kernel_fc12.py,megamoe_kernel.py,token_comm.py,sym_buffer.py,flag_batch.py}`	Workspace regions, pointer views, packed metadata, token-back modes, synchronization, release batching, and combine-format payload routing are reworked.
Epilogue pipelines and top-k reduction `tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/{epilogue_refactor.py,topk_reduce.py}`	FC1/FC2 epilogues use centralized quantization and routed combine stores, while device-side BF16, MXFP8, and FP4 reductions produce token-hidden outputs.
Backend bucketing and distributed setup `tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py`, `tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py`	Adaptive token buckets, per-bucket providers, MPI rendezvous, profiling scratch, and direct reduced outputs are wired into backend execution.
Streamed weight lifecycle `tensorrt_llm/_torch/modules/fused_moe/quantization.py`	Streamed NVFP4 source tensors are materialized during loading, checked for complete coverage, transformed into derived weights, and shrunk afterward.
Integration and validation `tensorrt_llm/llmapi/llm_args.py`, `tests/unittest/_torch/modules/moe/`, `tests/microbenchmarks/bench_moe/`	The new backend is registered across configuration and benchmark paths, and tests cover tactic selection, adaptive tuning, distributed setup, and streamed loading.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Scheduler
  participant MegaMoECuteDsl
  participant Sm100MegaMoEKernel
  participant TokenInPullTokenBackPush
  participant TopkReduce
  Scheduler->>MegaMoECuteDsl: provide rank-identical token count
  MegaMoECuteDsl->>MegaMoECuteDsl: select launch bucket and acquire buffers
  MegaMoECuteDsl->>Sm100MegaMoEKernel: launch fused dispatch and FC1/FC2
  Sm100MegaMoEKernel->>TokenInPullTokenBackPush: exchange token-back DATA and scale payloads
  Sm100MegaMoEKernel->>TopkReduce: reduce combine cells when separate reduction is selected
  TopkReduce-->>MegaMoECuteDsl: return token-hidden output

Suggested reviewers: mikeiovine, arysef, mzweilz, qijune, lfr-0531

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	No PR description was provided, so the required Description, Test Coverage, and PR Checklist sections are missing.	Add the template sections with a brief issue summary, the relevant tests run or added, and checklist items completed or noted.
Docstring Coverage	⚠️ Warning	Docstring coverage is 41.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise, follows the required template, and matches the PR's main MegaMoE CuTeDSL kernel updates.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (2)

tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py (1)
36-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Apply the required function annotations consistently across the new tests.

Both new test surfaces omit annotations required by the Python coding standard.

tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py#L36-L180: annotate helpers, nested helpers, _StreamingMoEModule.__init__, and test functions.

tests/unittest/_torch/modules/moe/test_moe_backend.py#L479-L525: annotate the constructor helper and all three test functions.

As per coding guidelines, “Annotate every function, use None for non-returning functions, avoid Any and unnecessary type ignores.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py` around
lines 36 - 180, Add explicit return and parameter annotations to every function
in tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py lines
36-180, including _load_classes, _StreamingMoEModule.__init__, nested
rand_u8/rand_fp8, _w13_input_scale, _w2_input_scale, _make_vanilla_weights,
_fresh, _load, _streamed_numels, _assert_sources_freed, _expected_mega_fc2,
_expected_mega_fc1, _expected_fc1_norm_const, and both test functions; use None
for non-returning functions and concrete container/tensor types without Any.
Apply the same annotations to the constructor helper and all three test
functions in tests/unittest/_torch/modules/moe/test_moe_backend.py lines
479-525.
Source: Coding guidelines
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py (1)
401-401: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a return type annotation to eval_function_mapping.

All return paths yield a dict, so annotate accordingly for consistency with the rest of the module.

As per coding guidelines: "Annotate every function, use None for non-returning functions".
Proposed change
-def eval_function_mapping(contract: Contract, **domain_coord):
+def eval_function_mapping(contract: Contract, **domain_coord) -> dict:
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py` at line 401,
Add a return type annotation of dict to eval_function_mapping, preserving its
existing behavior and all return paths.
Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/microbenchmarks/bench_moe/timing/autotune.py`:
- Around line 61-67: In the autotune flow around moe.forward() and
torch.cuda.synchronize(), determine the tactic_autotune flag before entering the
try block, then add conditional profiling-scratch release to the cleanup path so
it runs on both success and failure. Preserve the existing tuner-state
restoration in finally and reuse release_megamoe_profiling_scratch for enabled
backends.

---

Nitpick comments:
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py`:
- Line 401: Add a return type annotation of dict to eval_function_mapping,
preserving its existing behavior and all return paths.

In `@tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py`:
- Around line 36-180: Add explicit return and parameter annotations to every
function in tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
lines 36-180, including _load_classes, _StreamingMoEModule.__init__, nested
rand_u8/rand_fp8, _w13_input_scale, _w2_input_scale, _make_vanilla_weights,
_fresh, _load, _streamed_numels, _assert_sources_freed, _expected_mega_fc2,
_expected_mega_fc1, _expected_fc1_norm_const, and both test functions; use None
for non-returning functions and concrete container/tensor types without Any.
Apply the same annotations to the constructor helper and all three test
functions in tests/unittest/_torch/modules/moe/test_moe_backend.py lines
479-525.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2e4ad531-b0fe-43a0-b057-a5d2fce1260c

📥 Commits

Reviewing files that changed from the base of the PR and between 128d020 and ec63e67.

📒 Files selected for processing (33)

tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py
tensorrt_llm/_torch/models/modeling_deepseekv4.py
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
tensorrt_llm/_torch/modules/fused_moe/create_moe.py
tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py
tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
tensorrt_llm/_torch/modules/fused_moe/quantization.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/usage/llm_args_golden_manifest.json
tests/integration/test_lists/test-db/l0_b200.yml
tests/microbenchmarks/bench_moe/backend.py
tests/microbenchmarks/bench_moe/build.py
tests/microbenchmarks/bench_moe/search.py
tests/microbenchmarks/bench_moe/timing/autotune.py
tests/microbenchmarks/bench_moe/utils.py
tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
tests/unittest/_torch/modules/moe/test_moe_backend.py

tensorrt-cicd · 2026-07-22T05:44:08Z

PR_Github #60899 [ run ] triggered by Bot. Commit: ec63e67 Link to invocation

tensorrt-cicd · 2026-07-22T08:14:04Z

PR_Github #60899 [ run ] completed with state FAILURE. Commit: ec63e67
/LLM/main/L0_MergeRequest_PR pipeline #49167 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…Seek-V4 Add the MegaMoE-CuteDSL MoE backend (moe_config.backend=MEGAMOE_CUTEDSL): a pure cute.compile JIT integration of the MegaMoE NVFP4 swap-AB kernel (vendored under cute_dsl_kernels/mega_moe_nvfp4) for DeepSeek-V4 DEP serving on Blackwell. - Vendored kernel package: dispatch/combine token communication with NVLink sense-reversing barriers, swap-AB FC1/FC2 scheduling, NVFP4 epilogue with E8M0 scale conversion, adaptive grid sizing. - Custom op wrapper: workspace/fence lifecycle keyed on symmetric-buffer identity with storage-weakref liveness (a recycled allocation address can never skip the barrier-reset fence; capture-safe eviction, latched tuning workspaces, lazy profiling scratch released after autotuner warmup), adaptive max-tokens bucket ladder. - Tactic autotuning is an explicit opt-in (tactic_autotune backend constructor parameter, default OFF): serving always uses the deterministic heuristic tactic even under global autotuner warmup; the microbenchmark's --autotune enables the sweep, which then runs through the standard AutoTuner (tuning mode only). The curated reduced space is the only tuning space. - Cross-rank combine wire format is an explicit, validated backend constructor parameter (bf16 default; 32e4m3xe8m0 / 16e2m1xbf16 for benchmarking, threaded create_moe -> backend -> op and part of the workspace/provider cache keys; quantized formats use separate-reduce and disable the incompatible in-kernel reduce). No environment variables besides MEGAMOE_MPI_INIT_TORCH_DIST (opt-out of the process-level torch.distributed bootstrap). - Backend module wires the staged post-load hooks (transform_weights / cache_derived_state), free-port MPI bootstrap for multi-rank tuning, and fail-fast guards for unsupported configs (EPLB reload, pure-DEP requirement). - moe_config.backend gains the MEGAMOE_CUTEDSL literal (golden manifest regenerated); model_engine releases the profiling scratch at the end of autotuner warmup. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

… MegaMoE-CuteDSL - DSv4 model wiring: swiglu clamp forwarded in tensor form, fused-A / kv_b_proj atomic weight groups with module-state-keyed required sets, cross-RPC partial-group stash with covering-resend audit (bucket keys pre-registered per cycle, aborts union them into a tainted set that only a covering sweep clears), and a sticky abort-cycle protocol exposed as the duck-typed abort_reload_cycle hook. - attn_sink is a module-owned parameter (graph-stable storage created at construction, in-place refresh on load, value-preserving staged hooks) so checkpoint sinks survive CUDA-graph capture, dummy init and direct-transport loads; dummy weight init skips .attn_sink and memory-tagged allocation preserves constructor values. - Streaming source load for the MegaMoE-CuteDSL quant method: 0-element placeholders materialized per load, eager per-module finalize once coverage completes (CPU staging on partial reloads keeps the GPU peak order-independent), per-expert-row streamed coverage guard, and cycle-start purge of per-cycle quant transients (_RELOAD_TRANSIENT_ATTRS) so an aborted cycle's scales can never surface in the recovery cycle's finalize. - Reload correctness guards shared by Linear/MoE: per-unit reload-coverage debt recorded before pre-reload re-registration and consumed only by keys actually rewritten, with finalize refusing while units are outstanding; a side-effect-free check_reload_capability preflight refuses EPLB/non-partial-capable quant methods before the destructive walk; the update_weights exception boundary aborts the cycle (latch cleared, duck-typed model hook) and abort_update_weights lets a coordinator broadcast the same abort to every rank. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Self-contained kernel microbench over the public create_moe/forward path: synthetic NVFP4 weights, EP/DEP via mpirun, sweeps tokens-per-rank / combine wire format (passed explicitly to the backend constructor), --autotune opting into the MegaMoE tactic sweep through the standard AutoTuner (eager and --cuda-graph timing modes), free-port bootstrap and between-point workspace reset so repeated points are independent. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Treat the non-expert weights in DeepSeek V4 NVFP4 checkpoints as FP8 block-scale weights while retaining the per-layer NVFP4 expert configuration. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Barry-Delaney · 2026-07-23T02:01:43Z

/bot run

tensorrt-cicd · 2026-07-23T02:07:18Z

PR_Github #61149 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py (1)

662-686: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Duplicate sentinel WorkTileInfo construction — extract a shared helper.

The is_swap_ab sentinel-init block in create() (Lines 662-686) is duplicated verbatim in _gen_work_from_cluster_idx() (Lines 1258-1280). Both explicitly document that the field set must track ext.WorkTileInfo's serialized shape exactly — a future field change only applied to one copy would silently desync MLIR serialization.

♻️ Suggested consolidation

`@staticmethod`
def _default_work_tile(ext, params) -> "MoEWorkTileInfo":
    if const_expr(params.is_swap_ab):
        return ext.WorkTileInfo(
            expert_idx=Int32(WorkTileState.DONE),
            tile_m_idx=Int32(0),
            tile_n_idx=Int32(0),
            cumulative_data_physical_row=Int32(0),
            cumulative_sf_physical_row=Int32(0),
            cumulative_token_block_count=Int32(0),
            valid_tokens_in_cta_tile=Int32(0),
            phase_and_peek=Int32(BlockPhase.None_),
        )
    return ext.WorkTileInfo(
        expert_idx=Int32(WorkTileState.DONE),
        tile_m_idx=Int32(0),
        tile_n_idx=Int32(0),
        cumulative_data_physical_row=Int32(0),
        cumulative_sf_physical_row=Int32(0),
        cumulative_token_block_count=Int32(0),
        valid_tokens_in_cta_cluster_tile=Int32(0),
        phase_and_peek=Int32(BlockPhase.None_),
        fc1_counter_index=Int32(0),
    )

Then call self._default_work_tile(ext, params) from both create() and _gen_work_from_cluster_idx().

Also applies to: 1258-1280

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py`
around lines 662 - 686, Extract the duplicated sentinel WorkTileInfo
construction into a shared _default_work_tile helper that selects the exact
field set based on params.is_swap_ab. Replace both initialization blocks in
create() and _gen_work_from_cluster_idx() with calls to this helper, preserving
the serialized shapes and sentinel values unchanged.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py`:
- Around line 662-686: Extract the duplicated sentinel WorkTileInfo construction
into a shared _default_work_tile helper that selects the exact field set based
on params.is_swap_ab. Replace both initialization blocks in create() and
_gen_work_from_cluster_idx() with calls to this helper, preserving the
serialized shapes and sentinel values unchanged.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7a64c324-04bb-4756-b59c-604e208c2c20

📥 Commits

Reviewing files that changed from the base of the PR and between ec63e67 and a6aa781.

📒 Files selected for processing (25)

tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py
tensorrt_llm/_torch/models/modeling_deepseekv4.py
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
tensorrt_llm/_torch/modules/fused_moe/create_moe.py
tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py
tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
tensorrt_llm/_torch/modules/fused_moe/quantization.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/usage/llm_args_golden_manifest.json

🚧 Files skipped from review as they are similar to previous changes (18)

tensorrt_llm/usage/llm_args_golden_manifest.json
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/init.py
tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
tensorrt_llm/_torch/models/modeling_deepseekv4.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
tensorrt_llm/_torch/modules/fused_moe/quantization.py
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py

tensorrt-cicd · 2026-07-23T04:28:26Z

PR_Github #61149 [ run ] completed with state FAILURE. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49397 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Barry-Delaney · 2026-07-23T09:36:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T09:42:19Z

PR_Github #61269 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

tensorrt-cicd · 2026-07-23T12:52:24Z

PR_Github #61269 [ run ] completed with state FAILURE. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49504 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Barry-Delaney · 2026-07-23T13:04:43Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T13:17:04Z

PR_Github #61294 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

tensorrt-cicd · 2026-07-23T18:45:01Z

PR_Github #61294 [ run ] completed with state FAILURE. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49527 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Barry-Delaney · 2026-07-23T19:09:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T19:23:28Z

PR_Github #61374 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

tensorrt-cicd · 2026-07-23T21:28:47Z

PR_Github #61374 [ run ] completed with state SUCCESS. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49602 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Barry-Delaney self-assigned this Jul 9, 2026

Barry-Delaney force-pushed the user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt branch 11 times, most recently from d4cb10a to 5cded49 Compare July 17, 2026 04:39

Barry-Delaney added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 17, 2026

Barry-Delaney force-pushed the user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt branch from 06e6965 to 468c64a Compare July 17, 2026 09:57

Barry-Delaney marked this pull request as ready for review July 17, 2026 09:57

Barry-Delaney requested review from a team as code owners July 17, 2026 09:57

Barry-Delaney requested review from QiJune, arysef, bo-nv, lori-ren, nv-xtf and sunnyqgg July 17, 2026 09:57

coderabbitai Bot reviewed Jul 22, 2026

View reviewed changes

Comment thread tests/microbenchmarks/bench_moe/timing/autotune.py Outdated

Barry-Delaney added 10 commits July 23, 2026 10:01

[TRTLLM-12950][fix] narrow MegaMoE-CuteDSL integration scope

2e76a3e

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

[TRTLLM-12950][fix] isolate MegaMoE experimental controls

7211cbb

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

[TRTLLM-12950][test] narrow MegaMoE-CuteDSL test scope

b0d9194

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

[TRTLLM-12950][fix] suppress MegaMoE Bandit false positives

9470830

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

[TRTLLM-12950][test] remove redundant combine format tests

579dcdc

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

[None][fix] Resolve DeepSeek V4 mixed-precision base layers

682a8f2

Treat the non-expert weights in DeepSeek V4 NVFP4 checkpoints as FP8 block-scale weights while retaining the per-layer NVFP4 expert configuration. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

[TRTLLM-12950][fix] address MegaMoE review feedback

a6aa781

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Barry-Delaney force-pushed the user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt branch from 59cc034 to a6aa781 Compare July 23, 2026 02:01

coderabbitai Bot reviewed Jul 23, 2026

View reviewed changes

[TRTLLM-12950][fix] enforce MegaMoE group hint minimum

f047416

Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Shixiaowei02 approved these changes Jul 24, 2026

View reviewed changes

Uh oh!

Conversation

Barry-Delaney commented Jul 9, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Test list changes

Test code changes (outside test lists)

Uh oh!

coderabbitai Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

Barry-Delaney commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Barry-Delaney commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Barry-Delaney commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Barry-Delaney commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Barry-Delaney commented Jul 9, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading